package cas.ihep.spark.rdd;

import cas.ihep.spark.util.LocableObject;
import org.apache.spark.Partition;
import org.apache.spark.SparkContext;
import org.apache.spark.TaskContext;
import org.apache.spark.rdd.EmptyRDD;
import scala.collection.AbstractIterator;
import scala.collection.Iterator;
import scala.collection.JavaConversions;
import scala.collection.Seq;
import scala.reflect.ClassTag;

import java.util.Collection;
import java.util.Collections;
import java.util.NoSuchElementException;

public final class UserRDD<T> extends EmptyRDD<T> {
    private Partition[] partitions;
    public UserRDD(SparkContext sc, Collection<T> coll, ClassTag<T> evidence$1) {
        super(sc, evidence$1);
        partitions=new UserPartition[coll.size()];
        int i=0,rddId=id();
        java.util.Iterator<T> iter=coll.iterator();
        while(iter.hasNext()){
            T data=iter.next();
            partitions[i]=new UserPartition<>(data,rddId,i++);
        }
    }

    public Partition[] getPartitions(){
        return partitions;
    }

    @SuppressWarnings("unchecked")
    public Seq<String> getPreferredLocations(Partition p){
        if(p instanceof UserPartition){
            UserPartition up=(UserPartition)p;
            if(up.object instanceof LocableObject){
                return JavaConversions.asScalaBuffer(((LocableObject)up.object).getPreferredLocations());
            }else{
                return JavaConversions.asScalaBuffer(Collections.EMPTY_LIST);
            }
        }
        throw new RuntimeException("never happens");
    }

    @SuppressWarnings("unchecked")
    public Iterator<T> compute(Partition s, TaskContext context){
        if(s instanceof UserPartition){
            return new ComputeIterator<>(context,((UserPartition<T>) s).object);
        }else{
            throw new UnsupportedOperationException(s.getClass().getName()+" is not supported by UserRDD");
        }
    }

    private static class ComputeIterator<T> extends AbstractIterator<T> {
        T object;
        TaskContext context;
        ComputeIterator(TaskContext ctx,T obj) {
            object=obj;
            context=ctx;
        }
        public boolean hasNext(){
            context.killTaskIfInterrupted();
            return object!=null;
        }
        public T next(){
            if(object!=null){
                T ret=object;
                object=null;
                return ret;
            }
            throw new NoSuchElementException();
        }
    }

    private static class UserPartition<T> implements Partition{
        long rddId;
        int index;
        T object;
        UserPartition(T obj, long id, int i){
            object=obj;
            rddId=id;
            index=i;
        }
        @Override
        public int index() {
            return index;
        }

        public int hashCode(){
            return (int)(41 * (41 + rddId) + index);
        }

        public boolean equals(Object other){
            if(other instanceof UserPartition){
                return rddId==((UserPartition) other).rddId && index==((UserPartition) other).index;
            }
            return false;
        }
    }

}
